#/usr/bin/perl                                                                                                                           
use POSIX "fmod"; # float modulus
use List::Util qw[min]; # min function
# GLOBAL
$grid_size_x = 300;
$grid_size_y = 300;
$locitot = 99;
$vertical_scaling_factor = 1.5; # for hex grid is 0.75, we'll take modulus of 1.5
$biopsy_size_x = 1; # size of biopsy, 10x10 = 100 cells maximum
$biopsy_size_y = 1; # biopsies are 1mm by 2mm correspond to 10 crypts by 20 crypts
$distogram_output = 0; # Set to 1 to output distograms

# Biopsies
$number_of_biopsies_to_sample = 300;
$number_of_biopsies_schedule;
$number_of_biopsies_schedule[0] = $number_of_biopsies_to_sample;

# Crypts per biopsy
$number_of_crypts_per_biopsy_to_sample = 1;
$number_of_crypts_per_biopsy_schedule;
# Take x crypts from a biopsy
$number_of_crypts_per_biopsy_schedule[0] = $number_of_crypts_per_biopsy_to_sample;
# Alternatively, we can scehdule taking 1, 2, .., 20, 50, 100 crypts per biopsy as follows:
#for $i (0..19){
#$number_of_crypts_per_biopsy_schedule[$i] = $i+1;
#}
#$number_of_crypts_per_biopsy_schedule[20] = 50;
#$number_of_crypts_per_biopsy_schedule[21] = 100;
$number_of_replicates = 1; # how many times we repeat taking the same number of biopsies and number of crypts per biopsy

# GLOBAL hashes
%grid;
%patterns;
%indexx, %indexy; 


# Seed the random number generator
srand (0);


# Do simulation results file globbing
$outputdir = "/home/rumen/simulation1";
#$query = "K=0_P=0.25_MU=1e-06_NU=1e-05_RSM=1_B=0.02_D=0.016_NN=100_T=7300_H=300_W=300";
$query = "K=1.*T=14600";

# Get all sweeps that match the query
opendir(DIR, $outputdir);
@files = grep { /$query/ } readdir(DIR);
closedir(DIR);

foreach $f (@files) {
    print STDERR "Processing sweep $f\n";
    $dir = $outputdir."/".$f;    
    opendir(DIR, $dir);    
    @infiles = grep { /rep(\d+)\.log/ } readdir(DIR);
    closedir(DIR);
    
    # Contains all the rep runs number, e.g. if rep45, rep3, rep1 are present, rephash stores 45, 3, 1
    my %rephash = ();   
    foreach $i (@infiles) {
	@repnumber = split (/rep|\./, $i);
	$rephash{$repnumber[1]}=1;
    }

    foreach $run (keys %rephash){
	# The name of the run:
	compute_divergence_main($dir."/rep".$run);
	print STDERR "Sweep ".$dir."/rep".$run." done.\n";
    }
}


sub compute_divergence_main(){

    my ($runfn) = @_;

    my %grid_coords_x =();
    my %grid_coords_y =();
    open (in,"$runfn.crypts");
    while(<in>){
	chomp;
	@row=split(/\,/);
	$grid_coords_x{$row[0]}=$row[1];
	$grid_coords_y{$row[0]}=$row[2];
    }
    close in;

    %patterns=();

    # Load all patterns in memory
    open (in,"$runfn.patterns");
    while(<in>){
	chomp;
	@row=split(/\,/);
	$patterns{$row[0]}="";
	for my $loc (0..$locitot){
	    # Hash only the locations that are different from 
	    # the wild type "0" microsattelite length	    
	    # Shaka Zulu proposes hashing only the clone_id key
	    # and as value concatenate "locus,value,locus,value"
	    if($row[$loc+1]!=0){
		# Leading comma ",loc,val,loc,val"
		$patterns{$row[0]}=$patterns{$row[0]}.",".$loc.",".$row[$loc+1];
	    }
	}
    }
    close in;

    
    # Go through the log file of grid state
    %day = ();    
    %grid=();
    $firstday = 1;
    $logday = 1;
    open (i1,"$runfn.log");
    #open (out, ">$runfn.divergence.biopsy.sampling.log");
    open (out2, ">$runfn.divergence.biopsy.sampling.locations.log");
    if($distogram_output){open (out3, ">$runfn.biopsy.sampling.distogram.log");}
    while(<i1>){
	chomp;
	@row=split(/\,/);
	if($day{$row[2]}){ # If day exists in hash, then just add the new line of logfile to the grid
	    # Important - throw it as a double in the hash
	    $grid{$grid_coords_x{$row[0]}*1.0}{$grid_coords_y{$row[0]}*1.0}=$row[1];	
	} else { 	# If there's a new day in the log file, put the day in hash e.g. $day{"182"} = 1
	    # Upon a new day, execute
	    if($firstday>1){		
		# Compute divergence as a function of 
		# A. number of biopsies
		# B. number of cells sampled within a biopsy
		
		for my $replicate (1..$number_of_replicates){
		    for my $biopsyind (0..$#number_of_biopsies_schedule){
			for my $cryptind (0..$#number_of_crypts_per_biopsy_schedule){
			    # Number of crypts per biopsy
			    $crypt = $number_of_crypts_per_biopsy_schedule[$cryptind];
			    # Number of biopsies to sample
			    $biopsy = $number_of_biopsies_schedule[$biopsyind];
			    # Choose the locations of biopsies randomly, but make sure no two biopsies overlap
			    # Sample wifout replacement
			    my $biopsy_count = 0;
			    my %biopsy_coords = ();
			    while($biopsy_count<$biopsy){
				# Pick a biopsy
				%biopsyhash = ();
				%biopsyhash = draw_random_locations(1,0,$grid_size_x-1-$biopsy_size_x,0,$grid_size_y-1-$biopsy_size_y);
				# Get x, y of biopsy (hex coords)
				for my $kx (keys %biopsyhash){
				    for my $ky (keys %{$biopsyhash{$kx}}){
					$bx = $kx;
					$by = $ky;
				    }
				}   
				my ($bx, $by) = hex_to_abs_coords($bx, $by);
				$f = 1;
				# go through all coords of biopsies
				for my $i (1..$biopsy_count){
				    if(abs($bx-$biopsy_coords{$i}{"x"})<=$biopsy_size_x &&
				       abs($by-$biopsy_coords{$i}{"y"})<=$biopsy_size_y ){
					# failed, re-choose a biopsy coordinates
					$f = 0;
					last;
				    }
				}
				if($f==1){
				    # add biopsy
				    $biopsy_count++;
				    $biopsy_coords{$biopsy_count}{"x"}=$bx;
				    $biopsy_coords{$biopsy_count}{"y"}=$by;
				}
			    }
			    
			    # Evenly distribute crypts in biopsies
			    my %number_of_crypts_per_biopsy = ();
			    
			    # Sprinkle crypts in biopsies evenly
			    if($biopsy>=$crypt){
				for my $b (1..$biopsy){
				    $number_of_crypts_per_biopsy{$b} = $crypt;
				}
			    } else {
				for my $b (1..$biopsy){
				    $number_of_crypts_per_biopsy{$b} = int($crypt/$biopsy); 		   
				}
				for my $b (1..($crypt%$biopsy)){
				    $number_of_crypts_per_biopsy{$b} += 1;
				}		    
			    }
			    
			    %indexx=();
			    %indexy=();
			    my $i=0;
			    
			    # Go through each biopsy and draw crypts
			    for my $b (1..$biopsy){		    
				my %cryptshash = draw_random_locations($number_of_crypts_per_biopsy{$b},
								       $biopsy_coords{$b}{"x"}, $biopsy_coords{$b}{"x"}+$biopsy_size_x,
								       $biopsy_coords{$b}{"y"}, $biopsy_coords{$b}{"y"}+$biopsy_size_y);
				# Fill up the list of cells		   
				for my $kx (keys %cryptshash){
				    for my $ky (keys %{$cryptshash{$kx}}){
					$indexx{$i}=$kx;
					$indexy{$i}=$ky;
					$i++;	    
				    }
				}    		    
			    }
			    
			    # Compute divergence
			    my ($div,%divpair) = compute_divergence($i,\%grid,\%patterns,\%indexx,\%indexy);
			    $rounded = sprintf("%.6f", $div);	    
			    
			    # Log locations of picked cells
			    print out2 $logday.",$replicate,$biopsy,$crypt,$rounded";
			    for my $loc (0..$i-1){
				print out2 ",".$indexx{$loc}.",".$indexy{$loc};
			    }
			    print out2 "\n";
			    if($distogram_output){
				# Log the distograms
				print out3 $logday.",$replicate,$biopsy,$crypt";
				for my $comparison (keys %divpair){
				    my $genetic_distance = sprintf("%.6f", $divpair{$comparison}{"gd"});	    
				    my $spatial_distance = sprintf("%.1f", $divpair{$comparison}{"sd"});
				    print out3 ",".$spatial_distance.",".$genetic_distance;
				}
				print out3 "\n";		
			    }	    			    
			}
		    }
		}		
		
	    }
	    %grid=();
	    $day{$row[2]} = 1;
	    $logday = $row[2];
	    # Important - throw it as a double in the hash
	    $grid{$grid_coords_x{$row[0]}*1.0}{$grid_coords_y{$row[0]}*1.0}=$row[1];		
	    $firstday++;	    
	}	    
    }    
    close i1;  

    # After the last line of file has been read, re-do the analysis for the last day.
		# Compute divergence as a function of 
		# A. number of biopsies
		# B. number of cells sampled within a biopsy
		
		for my $replicate (1..$number_of_replicates){
		    for my $biopsyind (0..$#number_of_biopsies_schedule){
			for my $cryptind (0..$#number_of_crypts_per_biopsy_schedule){
			    # Number of crypts per biopsy
			    $crypt = $number_of_crypts_per_biopsy_schedule[$cryptind];
			    # Number of biopsies to sample
			    $biopsy = $number_of_biopsies_schedule[$biopsyind];
			    # Choose the locations of biopsies randomly, but make sure no two biopsies overlap
			    # Sample wifout replacement
			    my $biopsy_count = 0;
			    my %biopsy_coords = ();
			    while($biopsy_count<$biopsy){
				# Pick a biopsy
				%biopsyhash = ();
				%biopsyhash = draw_random_locations(1,0,$grid_size_x-1-$biopsy_size_x,0,$grid_size_y-1-$biopsy_size_y);
				# Get x, y of biopsy (hex coords)
				for my $kx (keys %biopsyhash){
				    for my $ky (keys %{$biopsyhash{$kx}}){
					$bx = $kx;
					$by = $ky;
				    }
				}   
				my ($bx, $by) = hex_to_abs_coords($bx, $by);
				$f = 1;
				# go through all coords of biopsies
				for my $i (1..$biopsy_count){
				    if(abs($bx-$biopsy_coords{$i}{"x"})<=$biopsy_size_x &&
				       abs($by-$biopsy_coords{$i}{"y"})<=$biopsy_size_y ){
					# failed, re-choose a biopsy coordinates
					$f = 0;
					last;
				    }
				}
				if($f==1){
				    # add biopsy
				    $biopsy_count++;
				    $biopsy_coords{$biopsy_count}{"x"}=$bx;
				    $biopsy_coords{$biopsy_count}{"y"}=$by;
				}
			    }
			    
			    # Evenly distribute crypts in biopsies
			    my %number_of_crypts_per_biopsy = ();
			    
			    # Sprinkle crypts in biopsies evenly
			    if($biopsy>=$crypt){
				for my $b (1..$biopsy){
				    $number_of_crypts_per_biopsy{$b} = $crypt;
				}
			    } else {
				for my $b (1..$biopsy){
				    $number_of_crypts_per_biopsy{$b} = int($crypt/$biopsy); 		   
				}
				for my $b (1..($crypt%$biopsy)){
				    $number_of_crypts_per_biopsy{$b} += 1;
				}		    
			    }
			    
			    %indexx=();
			    %indexy=();
			    my $i=0;
			    
			    # Go through each biopsy and draw crypts
			    for my $b (1..$biopsy){		    
				my %cryptshash = draw_random_locations($number_of_crypts_per_biopsy{$b},
								       $biopsy_coords{$b}{"x"}, $biopsy_coords{$b}{"x"}+$biopsy_size_x,
								       $biopsy_coords{$b}{"y"}, $biopsy_coords{$b}{"y"}+$biopsy_size_y);
				# Fill up the list of cells		   
				for my $kx (keys %cryptshash){
				    for my $ky (keys %{$cryptshash{$kx}}){
					$indexx{$i}=$kx;
					$indexy{$i}=$ky;
					$i++;	    
				    }
				}    		    
			    }
			    
			    # Compute divergence
			    my ($div,%divpair) = compute_divergence($i,\%grid,\%patterns,\%indexx,\%indexy);
			    $rounded = sprintf("%.6f", $div);	    			    
			    # Log locations of picked cells
			    print out2 $logday.",$replicate,$biopsy,$crypt,$rounded";
			    for my $loc (0..$i-1){
				print out2 ",".$indexx{$loc}.",".$indexy{$loc};
			    }
			    print out2 "\n";
			    if($distogram_output){
				# Log the distograms
				print out3 $logday.",$replicate,$biopsy,$crypt";
				for my $comparison (keys %divpair){
				    my $genetic_distance = sprintf("%.6f", $divpair{$comparison}{"gd"});	    
				my $spatial_distance = sprintf("%.1f", $divpair{$comparison}{"sd"});
				    print out3 ",".$spatial_distance.",".$genetic_distance;
				}
				print out3 "\n";			    
			    }
			}
		    }
		}		

  
    close out2;
    #close out;
    if($distogram_output){
	close out3;
    }
}


sub compute_divergence{
    # Number of cells
    my ($n,$gridRef,$patternsRef,$indexxRef,$indexyRef) = @_;
    my %grid = %$gridRef;
    my %patterns = %$patternsRef;
    my %indexx = %$indexxRef;
    my %indexy = %$indexyRef;
    my %divergenceofpair = ();

    my $divergence = 0.0;
    my $dead_cells_count = 0;

    my $comparison_count = 0;
    # Compare pattern of cell i to patterns of cell j
    # A total of n*(n-1)/2 comparisons 
    for my $i (0..($n-2)){
	my @col;
	# get the clone_id
	my $cid = $grid{$indexx{$i}}{$indexy{$i}};	

	# Skip comparing dead cells
	if($cid==-1){
	    $dead_cells_count++;
	}
	# count if last cell is dead
	if($i==$n-2 && $grid{$indexx{$i+1}}{$indexy{$i+1}}==-1){	
	    $dead_cells_count++;
	}

	if($cid!=-1){

	    for my $j (($i+1)..($n-1)){
		my $cid2 = $grid{$indexx{$j}}{$indexy{$j}};		

		# Speedup: if cid2 is a dead cell, skip it.
		if($cid2!=-1){
		    # Significant speedup: If the two chosen cells are the same clone
		    # then divergence is 0 between them, thus skip looping over
		    # their microsattelite patterns
		    # Only when they are different clones, compare their pattern
		    my $hamming = 0;
		    if($cid!=$cid2){
						
			# Convert pattern
			my @cell_one = split(/\,/,$patterns{$cid});
			my @cell_two = split(/\,/,$patterns{$cid2});
			
			my %h1 = ();
			my %h2 = ();

			for my $crnch (1..$#cell_one){
			    if($crnch%2==1){
				# Hash locus, value
				# If value is greater than 2, set to 2
				# MAJOR CHANGE: Simulate 0-no loss 1-single copy loss 2-double copy loss states
				# Versus 0-infinity neutral states
				$h1{$cell_one[$crnch]} = ($cell_one[$crnch+1]>2)?2:$cell_one[$crnch+1];
			    }
			}
			for my $crnch (1..$#cell_two){
			    if($crnch%2==1){
				# Hash locus, value
				$h2{$cell_two[$crnch]} = ($cell_two[$crnch+1]>2)?2:$cell_two[$crnch+1];
			    }
			}

			# Loop over cell one hashed pattern
			while (($locus, $value) = each(%h1)) {
			    #print "$cid $locus $value";
			    # If the locus is present in cell two
			    if(defined($h2{$locus})){
				$hamming += abs($h1{$locus}-$h2{$locus});			    
				#$hamming += abs(min(2,$h1{$locus})-min(2,$h2{$locus}));				
			    } else {
				$hamming += abs($h1{$locus});
				#$hamming += abs(min(2,$h1{$locus}));
			    }
			}
			#print "Cell TWO:\n";
			# Loop over cell two hashed pattern
			while (($locus, $value) = each(%h2)) {
			    #print "$cid2 $locus $value";
			    # If the locus is present in cell one, do nothing, if absent, compare to 100
			    if(!defined($h1{$locus})){			    
				$hamming += abs($h2{$locus});
				#$hamming += abs(min(2,$h2{$locus}));			 
			    }
			}
			# The rest of the loci are all 100s, so hamming distance is 0.
			# The distance between two cells is hamming/#of valid comparisons
			$divergence += $hamming*1.0/($locitot+1);
			
		    }
		    if($distogram_output){
			# Save divergence of the pair of cells whether or not they are from the same clone
			# Set divergence of the chosen two cells to -1 (comparing dead cells)
			$divergenceofpair{$comparison_count}{"gd"} = $hamming*1.0/($locitot+1);
			my ($ax, $ay) = hex_to_abs_coords($indexx{$i}, $indexy{$i});
			my ($bx, $by) = hex_to_abs_coords($indexx{$j}, $indexy{$j});
			# Distance base on wrap around hex grid
			my $xd = min ($grid_size_x-abs($ax-$bx), abs($ax-$bx));
			my $yd = abs($ay-$by);
			$divergenceofpair{$comparison_count}{"sd"} = sqrt($xd*$xd+$yd*$yd);
		    }
		    $comparison_count++;
		}
	    }
	}
	#print "Comparing cell $i\n";
    }
    #print $n."\n";
    # At the end divide by the total comparisons between cells
    # subtract dead cells count from computation
    # What if all cells that we picked are dead? Then, set divergence to 0.
    if($dead_cells_count<$n-1){
	$divergence = $divergence/(($n-$dead_cells_count)*($n-$dead_cells_count-1)/2.0);
    } else {
	$divergence = 0.0;
    }
    return ($divergence,%divergenceofpair);
}


# Draws random n locations
sub draw_random_locations{
    my ($n,$leftx,$rightx,$topy,$boty) = @_;
    my %h=();
    my $size=0;
    while($size<$n){
        $x = $leftx+int(rand($rightx-$leftx+1)); 
        $y = $topy+int(rand($boty-$topy+1)); 	
	# Convert abs to hex
	my ($hx, $hy) = abs_to_hex_coords($x, $y);
	# Ensure we pick out n DIFFERENT cells
	# do not pick out the same cell twice
        if(!$h{$hx}{$hy}){
            $h{$hx}{$hy}=1;
	    $size++;
        }
    }
    return %h;
}

# convert coords
sub abs_to_hex_coords{
    my ($x, $y) = @_;    
    $y = $y*($vertical_scaling_factor/2.0);
    # Depending on picking y that is divisible by 1.5 or not, add 0.5 to x                                         
    if(fmod($y,$vertical_scaling_factor)){
	$x = $x + 0.5;
    }
    return ($x, $y);
}

sub hex_to_abs_coords{
    my ($x, $y) = @_;    
    # Depending on picking y that is divisible by 1.5 or not, add 0.5 to x                                         
    if(fmod($y,$vertical_scaling_factor)){
	$x = $x - 0.5;
    }
    $y = $y/($vertical_scaling_factor/2.0);
    return ($x, $y);
}

